# coding:utf-8
import numpy as np
import cv2
import matplotlib.pyplot as plt
from PIL import Image
import glob
#-------------------------------
# 设置棋盘格大小  Checkerboard参数
#
#-------------------------------
Checkerboard = (9,6)  #棋盘格内角点数
square_size = (10,10)   #棋盘格大小，单位mm

# 设置迭代终止条件
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)                 # 单目的，亚像素函数使用
criteria_stereo = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)           # 双目的。亚像素函数使用


# 设置 object points, 形式为 (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((Checkerboard[0] * Checkerboard[1], 3), np.float32)  # 我用的是6×7的棋盘格，可根据自己棋盘格自行修改相关参数
objp[:, :2] = np.mgrid[0:Checkerboard[0], 0:Checkerboard[1]].T.reshape(-1, 2)

# 用arrays存储所有图片的object points 和 image points
objpoints = []  # 3d points in real world space
imgpointsR = []  # 2d points in image plane
imgpointsL = []

imagL = glob.glob(r"./img/left/*.jpg")
imagR = glob.glob(r"./img/right/*.jpg")
# 本次实验采集里共计30组待标定图片依次读入进行以下操作
for i in range(len(imagL)):
    t = str(i)
    ImaL = cv2.imread(imagL[i], 1)  # 右视图
    ImaR = cv2.imread(imagR[i], 1)  # 左视图
    grayL = cv2.cvtColor(ImaL, cv2.COLOR_BGR2GRAY)
    grayR = cv2.cvtColor(ImaR, cv2.COLOR_BGR2GRAY)
    retR, cornersR = cv2.findChessboardCorners(grayR, Checkerboard, None)  # 提取右图每一张图片的角点
    retL, cornersL = cv2.findChessboardCorners(grayL, Checkerboard, None)  # # 提取左图每一张图片的角点
    if (True == retR) & (True == retL):
        objpoints.append(objp)
        cv2.cornerSubPix(grayL, cornersL, (11, 11), (-1, -1), criteria)  # 亚像素精确化，对粗提取的角点进行精确化
        cv2.cornerSubPix(grayR, cornersR, (11, 11), (-1, -1), criteria)  # 亚像素精确化，对粗提取的角点进行精确化
        imgpointsR.append(cornersR)
        imgpointsL.append(cornersL)

        cv2.drawChessboardCorners(ImaL, Checkerboard, cornersL, retL)
        cv2.drawChessboardCorners(ImaR, Checkerboard, cornersR, retR)

        # 把数据压到栈中,一起显示左右检测结果
        Image_merge = np.hstack((ImaL, ImaR))
        # 各参数依次是：图片，添加的文字，左上角坐标，字体，字体大小，颜色，字体粗细
        cv2.putText(Image_merge,"Left_img", (5, 40), 0, 1, (0, 0, 255), 2)
        cv2.putText(Image_merge, "Right_img", (5+ImaL.shape[1], 40), 0, 1, (0, 0, 255), 2)

        cv2.imshow('input img', Image_merge)
        print("请敲击任意键，检测下一帧角点...")
        cv2.waitKey(100)
    else:
        print("检测角点失败...,请检查图片路径或者棋盘格大小")
cv2.destroyAllWindows()
# 相机的单双目标定、及校正
#   右侧相机单独标定
retR, mtxR, distR, rvecsR, tvecsR = cv2.calibrateCamera(objpoints, imgpointsR, grayR.shape[::-1], None, None)
print("右相机内参：\n",mtxR, "右相机畸变：\n",distR)
print("\n","-"*30,"\n")
#   获取新的相机矩阵后续传递给initUndistortRectifyMap，以用remap生成映射关系
hR, wR = grayR.shape[:2]
OmtxR, roiR = cv2.getOptimalNewCameraMatrix(mtxR, distR, (wR, hR), 1, (wR, hR))

#   左侧相机单独标定
retL, mtxL, distL, rvecsL, tvecsL = cv2.calibrateCamera(objpoints, imgpointsL, grayL.shape[::-1], None, None)
print("左相机内参：\n",mtxL, "左相机畸变：\n",distL)
print("\n","-"*30,"\n")
#   获取新的相机矩阵后续传递给initUndistortRectifyMap，以用remap生成映射关系
hL, wL = grayL.shape[:2]
OmtxL, roiL = cv2.getOptimalNewCameraMatrix(mtxL, distL, (wL, hL), 1, (wL, hL))

# 双目相机的标定
# 设置标志位为cv2.CALIB_FIX_INTRINSIC，这样就会固定输入的cameraMatrix和distCoeffs不变，只求解𝑅,𝑇,𝐸,𝐹
flags = 0
flags |= cv2.CALIB_FIX_INTRINSIC

retS, MLS, dLS, MRS, dRS, R, T, E, F = cv2.stereoCalibrate(objpoints, imgpointsL, imgpointsR, OmtxL, distL, OmtxR,
                                                            distR,
                                                           grayR.shape[::-1], criteria_stereo, flags)
print("左右相机外参-旋转分量：\n",R,"\n左右相机外参-平移分量：\n",T)
print("\n","-"*30,"\n")

# 利用stereoRectify()计算立体校正的映射矩阵
rectify_scale = 1  # 设置为0的话，对图片进行剪裁，设置为1则保留所有原图像像素
RL, RR, PL, PR, Q, roiL, roiR = cv2.stereoRectify(MLS, dLS, MRS, dRS,
                                                  grayR.shape[::-1], R, T,
                                                  rectify_scale, (0, 0))
# 利用initUndistortRectifyMap函数计算畸变矫正和立体校正的映射变换，实现极线对齐。
Left_Stereo_Map = cv2.initUndistortRectifyMap(MLS, dLS, RL, PL,
                                              grayR.shape[::-1], cv2.CV_16SC2)

Right_Stereo_Map = cv2.initUndistortRectifyMap(MRS, dRS, RR, PR,
                                               grayR.shape[::-1], cv2.CV_16SC2)


# 立体校正效果显示
count_img = 14
for i in range(count_img):  # 以第一对图片为例
    t = str(i)
    frameR = cv2.imread(imagL[i], 0)
    frameL = cv2.imread(imagR[i], 0)

    Left_rectified = cv2.remap(frameL, Left_Stereo_Map[0], Left_Stereo_Map[1], cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT,
                               0)  # 使用remap函数完成映射
    im_L = Image.fromarray(Left_rectified)  # numpy 转 image类

    Right_rectified = cv2.remap(frameR, Right_Stereo_Map[0], Right_Stereo_Map[1], cv2.INTER_LANCZOS4,
                                cv2.BORDER_CONSTANT, 0)
    im_R = Image.fromarray(Right_rectified)  # numpy 转 image 类

    Image_rectified = np.hstack((Left_rectified, Right_rectified))
    # 各参数依次是：图片，添加的文字，左上角坐标，字体，字体大小，颜色，字体粗细
    cv2.putText(Image_rectified, "Left_img", (5, 40), 0, 1, (0, 0, 255), 2)
    cv2.putText(Image_rectified, "Right_img", (5 + Left_rectified.shape[1], 40), 0, 1, (0, 0, 255), 2)

    cv2.imshow('input img', Image_rectified)
    print("请敲击任意键，检测下一帧角点...")
    cv2.waitKey()
    print("按任意键显示下一对矫正图像...")

    # 创建一个能同时并排放下两张图片的区域，后把两张图片依次粘贴进去
    width = im_L.size[0] * 2
    height = im_L.size[1]
    img_compare = Image.new('RGBA', (width, height))
    img_compare.paste(im_L, box=(0, 0))
    img_compare.paste(im_R, box=(im_L.size[0], 0))
    # 在已经极线对齐的图片上均匀画线
    for i in range(1, 20):
        len = im_L.size[1] / 20
        plt.axhline(y=i * len, color='r', linestyle='-')
    plt.imshow(img_compare)
    plt.show()
cv2.destroyAllWindows()

img1 = cv2.imread(imagL[0])
img2 = cv2.imread(imagR[0])

cv2.namedWindow("left")
cv2.namedWindow("right")
cv2.namedWindow("depth")
cv2.resizeWindow("depth", 1200, 800);
cv2.moveWindow("left", 0, 0)
cv2.moveWindow("right", 600, 0)
cv2.createTrackbar("num", "depth", 0, 10, lambda x: None)
cv2.createTrackbar("blockSize", "depth", 5, 255, lambda x: None)


# 添加点击事件，打印当前点的距离
def callbackFunc(e, x, y, f, p):
    if e == cv2.EVENT_LBUTTONDOWN:
        print(threeD[y][x])


cv2.setMouseCallback("depth", callbackFunc, None)

Left_rectified = cv2.remap(img1, Left_Stereo_Map[0], Left_Stereo_Map[1], cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT, 0)
Right_rectified = cv2.remap(img2, Right_Stereo_Map[0], Right_Stereo_Map[1], cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT, 0)

# 将图片置为灰度图，为StereoBM作准备
# cv2.imshow("gray", Left_rectified)
imgL = cv2.cvtColor(Left_rectified, cv2.COLOR_BGR2GRAY)
imgR = cv2.cvtColor(Right_rectified, cv2.COLOR_BGR2GRAY)

# 两个trackbar用来调节不同的参数查看效果
num = cv2.getTrackbarPos("num", "depth")
blockSize = cv2.getTrackbarPos("blockSize", "depth")
if blockSize % 2 == 0:
    blockSize += 1
if blockSize < 5:
    blockSize = 5

# 根据Block Maching方法生成差异图（opencv里也提供了SGBM/Semi-Global Block Matching算法，有兴趣可以试试）
stereo = cv2.StereoBM_create(numDisparities=16 * num, blockSize=blockSize)
disparity = stereo.compute(imgL, imgR)

disp = cv2.normalize(disparity, disparity, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
# 将图片扩展至3d空间中，其z方向的值则为当前的距离
threeD = cv2.reprojectImageTo3D(disparity.astype(np.float32) / 16., Q)

cv2.imshow("left", imgL)
cv2.imshow("right", imgR)
cv2.imshow("depth", disp)
cv2.imwrite("./depth.bmp", disp)

key = cv2.waitKey(0)
if key == ord("q"):
    cv2.destroyAllWindows()